#!/usr/bin/env python3

import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import cv2
import numpy as np
import serial
import time
import math


class StereoBallDetector(Node):
    def __init__(self):
        super().__init__('stereo_ball_detector')
        
        # -------------------------- 常量与参数初始化 --------------------------
        # 协议常量（与STM32匹配）
        self.FRAME_HEADER = 0x66
        self.FRAME_LENGTH = 0x07
        self.FRAME_TAIL = 0x5B
        self.SEND_INTERVAL = 0.05  # 发送间隔（20Hz）
        
        # 初始化相机参数（双目矫正参数）
        self.cam_params = self.init_camera_params()
        self.screen_center = (self.cam_params["size"][0] // 2, self.cam_params["size"][1] // 2)
        self.get_logger().info(f"图像中心点: {self.screen_center}")
        
        # 初始化立体匹配器
        self.stereo, self.right_matcher, self.wls_filter = self.init_stereo_matcher()
        
        # 初始化串口
        self.ser = self.init_serial()
        
        # 图像处理工具
        self.bridge = CvBridge()
        
        # 时间控制
        self.last_send_time = time.time()
        
        # 订阅usb_cam发布的图像（1280x480，左右拼接）
        self.image_sub = self.create_subscription(
            Image,
            '/image_raw',  # usb_cam默认话题
            self.image_callback,
            10
        )
        
        # 创建显示窗口
        cv2.namedWindow("Rectified Color", cv2.WINDOW_NORMAL)
        cv2.namedWindow("Depth Map", cv2.WINDOW_NORMAL)
        cv2.setMouseCallback("Depth Map", self.onmouse_pick_points)
        
        self.get_logger().info("双目橙色球检测节点已启动")

    # -------------------------- 初始化函数 --------------------------
    def init_camera_params(self):
        """初始化双目相机参数（包含畸变矫正参数）"""
        left_camera_matrix = np.array([[715.088662856028, -1.27750597113329, 321.725929157012],
                                       [0, 717.938092436680, 251.352481762661],
                                       [0., 0., 1.]])
        right_camera_matrix = np.array([[712.537452801351, -1.49459241948369, 325.232106968860],
                                        [0, 715.552769066380, 250.629910089625],
                                        [0., 0., 1.]])

        left_distortion = np.array(
            [[-0.440986151974959, 0.327464329751247, -0.000547150225973162, -0.00126940700042557, -0.335074894313998]])
        right_distortion = np.array(
            [[-0.444122749137210, 0.348225581098782, 0.00176792559350367, 0.00123611244757822, -0.660993977721134]])

        R = np.array([[0.999866011965132, 0.000416637643297468, 0.0163641232584165],
                      [-0.000544744807817963, 0.999969236927380, 0.00782486114692518],
                      [-0.0163603597159969, -0.00783272698033651, 0.999835479975588]])
        T = np.array([-59.8387197927031, 0.0107452720401373, 0.330883446596301])
        size = (640, 480)

        # 计算立体校正参数
        R1, R2, P1, P2, Q, validPixROI1, validPixROI2 = cv2.stereoRectify(
            left_camera_matrix, left_distortion,
            right_camera_matrix, right_distortion,
            size, R, T
        )

        # 生成左右相机的矫正映射表
        left_map1, left_map2 = cv2.initUndistortRectifyMap(
            left_camera_matrix, left_distortion, R1, P1, size, cv2.CV_16SC2)
        right_map1, right_map2 = cv2.initUndistortRectifyMap(
            right_camera_matrix, right_distortion, R2, P2, size, cv2.CV_16SC2)

        return {
            "left_camera_matrix": left_camera_matrix,
            "right_camera_matrix": right_camera_matrix,
            "left_distortion": left_distortion,
            "right_distortion": right_distortion,
            "R": R,
            "T": T,
            "size": size,
            "Q": Q,
            "left_map1": left_map1,
            "left_map2": left_map2,
            "right_map1": right_map1,
            "right_map2": right_map2,
            "baseline": abs(T[0])
        }

    def init_stereo_matcher(self, blockSize=5, numDisparities=128):
        """初始化SGBM匹配器和WLS滤波器"""
        img_channels = 3
        stereo = cv2.StereoSGBM_create(
            minDisparity=0,
            numDisparities=numDisparities,
            blockSize=blockSize,
            P1=8 * img_channels * blockSize * blockSize,
            P2=32 * img_channels * blockSize * blockSize,
            disp12MaxDiff=5,
            preFilterCap=63,
            uniquenessRatio=15,
            speckleWindowSize=200,
            speckleRange=32,
            mode=cv2.STEREO_SGBM_MODE_SGBM_3WAY
        )

        right_matcher = cv2.ximgproc.createRightMatcher(stereo)
        wls_filter = cv2.ximgproc.createDisparityWLSFilter(matcher_left=stereo)
        wls_filter.setLambda(8000)
        wls_filter.setSigmaColor(1.5)

        return stereo, right_matcher, wls_filter

    def init_serial(self, port='/dev/ttyUSB0', baudrate=115200):
        """初始化串口通信（与STM32通信）"""
        try:
            ser = serial.Serial(
                port=port,
                baudrate=baudrate,
                bytesize=serial.EIGHTBITS,
                parity=serial.PARITY_NONE,
                stopbits=serial.STOPBITS_ONE,
                timeout=0.1
            )
            self.get_logger().info(f"串口已打开: {ser.port}")
            return ser
        except serial.SerialException as e:
            self.get_logger().error(f"串口打开失败: {e}")
            self.get_logger().warn("将以无串口模式运行")
            return None

    # -------------------------- 图像处理函数 --------------------------
    def rectify_image(self, img, map1, map2):
        """对图像进行畸变矫正"""
        return cv2.remap(img, map1, map2, cv2.INTER_LINEAR)

    def compute_disparity_and_depth(self, left_rect, right_rect):
        """计算视差图和深度信息"""
        disparity_left = self.stereo.compute(left_rect, right_rect)
        disparity_right = self.right_matcher.compute(right_rect, left_rect)

        filtered_disparity = self.wls_filter.filter(disparity_left, left_rect, None, disparity_right)
        filtered_disparity = np.clip(filtered_disparity, 0, None)

        # 视差图归一化（用于显示）
        disp_norm = cv2.normalize(filtered_disparity, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U)
        disp_color = cv2.applyColorMap(disp_norm, cv2.COLORMAP_JET)

        # 计算三维坐标（深度信息）
        threeD = cv2.reprojectImageTo3D(filtered_disparity, self.cam_params["Q"], handleMissingValues=False)
        threeD = threeD * 16.0  # 缩放深度值

        return filtered_disparity, disp_norm, disp_color, threeD

    def detect_orange_ball(self, rect_color, threeD):
        """在矫正后的彩色图上检测橙色球"""
        # 转换到HSV颜色空间
        hsv = cv2.cvtColor(rect_color, cv2.COLOR_BGR2HSV)
        lower_orange = np.array([15, 82, 210])
        upper_orange = np.array([68, 255, 255])
        mask = cv2.inRange(hsv, lower_orange, upper_orange)

        # 形态学操作去除噪声
        kernel = np.ones((5, 5), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)

        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        ball_info = None
        best_score = 0

        for contour in contours:
            area = cv2.contourArea(contour)
            if area < 50:
                continue

            # 计算圆形度
            perimeter = cv2.arcLength(contour, True)
            if perimeter == 0:
                continue
            circularity = 4 * math.pi * area / (perimeter **2)
            if circularity < 0.6:
                continue

            # 计算轮廓中心（矫正后坐标）
            (cx_rect, cy_rect), radius = cv2.minEnclosingCircle(contour)
            cx_rect, cy_rect = int(cx_rect), int(cy_rect)

            # 验证深度有效性
            if 0 <= cx_rect < threeD.shape[1] and 0 <= cy_rect < threeD.shape[0]:
                depth = threeD[cy_rect, cx_rect]
                z_val = depth[2]
                if z_val <= 0:
                    continue

                # 计算实际距离（米）
                distance = np.linalg.norm(depth) / 1000.0

                # 评分（优先选择面积大、圆形度高的目标）
                score = area * circularity
                if score > best_score:
                    best_score = score
                    ball_info = {
                        "rect_center": (cx_rect, cy_rect),
                        "depth": depth,
                        "distance": distance,
                        "contour": contour
                    }

        return ball_info

    # -------------------------- 通信与交互函数 --------------------------
    def send_coordinate(self, x, y):
        """按协议发送坐标给STM32"""
        if not self.ser or not self.ser.is_open:
            return

        # 限制坐标范围为0~65535
        x_clamped = np.clip(int(x), 0, 65535)
        y_clamped = np.clip(int(y), 0, 65535)

        # 拆分高低字节
        x_low = x_clamped & 0xFF
        x_high = (x_clamped >> 8) & 0xFF
        y_low = y_clamped & 0xFF
        y_high = (y_clamped >> 8) & 0xFF

        # 构建帧结构
        frame = bytearray([
            self.FRAME_HEADER, self.FRAME_LENGTH,
            x_low, x_high, y_low, y_high,
            self.FRAME_TAIL
        ])

        # 发送数据
        try:
            self.ser.write(frame)
            self.get_logger().debug(f"发送坐标: ({x_clamped}, {y_clamped})  帧: {[hex(b) for b in frame]}")
        except serial.SerialException as e:
            self.get_logger().error(f"串口发送失败: {e}")

    def onmouse_pick_points(self, event, x, y, flags, param):
        """鼠标回调：点击深度图获取三维坐标（调试用）"""
        if hasattr(self, 'threeD') and event == cv2.EVENT_LBUTTONDOWN:
            threeD = self.threeD
            if 0 <= y < threeD.shape[0] and 0 <= x < threeD.shape[1]:
                depth = threeD[y, x]
                z_val = depth[2]
                if z_val <= 0:
                    self.get_logger().info("无效深度值")
                    return
                xyz = depth / 1000.0  # 转换为米
                distance = np.linalg.norm(xyz)
                self.get_logger().info(
                    f"点击坐标: ({x}, {y})  三维坐标: ({xyz[0]:.3f}, {xyz[1]:.3f}, {xyz[2]:.3f})m  距离: {distance:.3f}m"
                )

    # -------------------------- 图像回调函数 --------------------------
    def image_callback(self, msg):
        """处理订阅的图像消息"""
        try:
            # 将ROS图像消息转换为OpenCV格式（BGR8）
            cv_image = self.bridge.imgmsg_to_cv2(msg, "bgr8")
        except Exception as e:
            self.get_logger().error(f"图像转换失败: {e}")
            return

        # 分割左右相机图像（1280x480 -> 左右各640x480）
        frame_left = cv_image[:, :640]  # 左相机原始彩色图
        frame_right = cv_image[:, 640:]  # 右相机原始彩色图

        # 转换为灰度图（用于立体匹配）
        gray_left = cv2.cvtColor(frame_left, cv2.COLOR_BGR2GRAY)
        gray_right = cv2.cvtColor(frame_right, cv2.COLOR_BGR2GRAY)

        # 畸变矫正
        left_rect_gray = self.rectify_image(gray_left, self.cam_params["left_map1"], self.cam_params["left_map2"])
        left_rect_color = self.rectify_image(frame_left, self.cam_params["left_map1"], self.cam_params["left_map2"])
        right_rect_gray = self.rectify_image(gray_right, self.cam_params["right_map1"], self.cam_params["right_map2"])

        # 计算视差和深度
        _, _, disp_color, threeD = self.compute_disparity_and_depth(left_rect_gray, right_rect_gray)
        self.threeD = threeD  # 保存深度信息供鼠标回调使用

        # 检测橙色球
        ball_info = self.detect_orange_ball(left_rect_color, threeD)

        # 绘制检测结果
        target_x, target_y = self.screen_center  # 默认发送中心点
        if ball_info:
            cx_rect, cy_rect = ball_info["rect_center"]
            distance = ball_info["distance"]
            contour = ball_info["contour"]
            target_x, target_y = cx_rect, cy_rect  # 更新为小球中心

            # 绘制轮廓和中心
            cv2.drawContours(left_rect_color, [contour], -1, (0, 255, 0), 2)
            cv2.circle(left_rect_color, (cx_rect, cy_rect), 5, (0, 0, 255), -1)
            # 绘制与中心点连线
            cv2.line(left_rect_color, self.screen_center, (cx_rect, cy_rect), (0, 255, 255), 2)
            # 显示信息
            cv2.putText(left_rect_color, f"Center: ({cx_rect}, {cy_rect})", (10, 30),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)
            cv2.putText(left_rect_color, f"Distance: {distance:.3f}m", (10, 60),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)
        else:
            cv2.putText(left_rect_color, "No Target", (10, 30),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)

        # 绘制图像中心点
        cv2.circle(left_rect_color, self.screen_center, 5, (255, 0, 0), -1)
        cv2.line(left_rect_color, (self.screen_center[0]-10, self.screen_center[1]),
                (self.screen_center[0]+10, self.screen_center[1]), (255, 0, 0), 2)
        cv2.line(left_rect_color, (self.screen_center[0], self.screen_center[1]-10),
                (self.screen_center[0], self.screen_center[1]+10), (255, 0, 0), 2)

        # 按间隔发送坐标
        current_time = time.time()
        if current_time - self.last_send_time > self.SEND_INTERVAL:
            self.send_coordinate(target_x, target_y)
            self.last_send_time = current_time

        # 显示图像
        cv2.imshow("Rectified Color", left_rect_color)
        cv2.imshow("Depth Map", disp_color)
        cv2.waitKey(1)  # 处理窗口事件


def main(args=None):
    rclpy.init(args=args)
    detector = StereoBallDetector()
    try:
        rclpy.spin(detector)
    except KeyboardInterrupt:
        pass
    finally:
        # 释放资源
        cv2.destroyAllWindows()
        if detector.ser and detector.ser.is_open:
            detector.ser.close()
            detector.get_logger().info("串口已关闭")
        detector.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
